import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GATConv, global_mean_pool


class PUTraceAD(nn.Module):
    def __init__(self, args):
        super(PUTraceAD, self).__init__()

        self.base_dir = args.base_dir
        self.model_fname = args.model_fname

        self.hidden = args.gnn_hidden

        self.gnn1 = GATConv(1, self.hidden, 3)
        self.gnn2 = GATConv(self.hidden*3, self.hidden, 3)
        self.gnn3 = GATConv(self.hidden*3, self.hidden, 1)
        self.fc = nn.Linear(self.hidden, 2)

    def forward(self, data):
        # in.x.shape: (~B*L, F_s)
        # out.shape: (B, 2)

        x, edge_index, batch = data.x, data.edge_index, data.batch
        x = self.gnn1(x, edge_index)
        x = self.gnn2(x, edge_index)
        x = self.gnn3(x, edge_index)

        # mean pooling
        vg = global_mean_pool(x, batch)   # (B, H)
        out = F.softmax(self.fc(vg))      # (B, 2)

        return out

    def load(self):
        file_path = os.path.join(self.base_dir, self.model_fname)
        self.load_state_dict(torch.load(file_path))

    def save(self):
        file_path = os.path.join(self.base_dir, self.model_fname)
        if os.path.exists(file_path):
            os.remove(file_path)
        torch.save(self.state_dict(), file_path)
